# $Id: snmp.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#
# essential imports
__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

from pyasn1.codec.ber import encoder, decoder
from pysnmp.carrier.asynsock.dispatch import AsynsockDispatcher
from pysnmp.carrier.asynsock.dgram import udp
from pysnmp.entity import engine, config
from pysnmp.entity.rfc3413 import cmdgen
from pysnmp.proto import api, rfc1902
from pysnmp.smi import builder, view, error
from time import time

import stdlib
from exception import SNMPError

GET = 1
SET = 2
WALK = 3

snmptype = stdlib.EnumType('GET', 'SET', 'WALK', 'BULKWALK')

class snmp(object):
   def __init__(self, **Keywords):
      self._host = None
      self._comm = None
      self._port = 0
      self._snmp = None
      self._type = None
      self._reqmesg = None
      self._conn = False

      if "server" in Keywords:
         self._host = Keywords["server"].strip()

      if "community" in Keywords:
         self._comm = Keywords["community"].strip()

      if "port" in Keywords:
         try:
            self._port = int(Keywords["port"])
         except ValueError, message:
            raise ValueError, message

      if self._port == 0:
         self._port = 161

      if "version" in Keywords:
         if Keywords["version"] == "1":
            self._snmp = api.protoModules[api.protoVersion1]
         elif Keywords["version"] == "2c":
            self._snmp = api.protoModules[api.protoVersion2c]
         elif Keywords["version"] == "2":
            self._snmp = api.protoModules[api.protoVersion2c]

      if not self._snmp:
         self._snmp = api.protoModules[api.protoVersion2c]

      self._builder = builder.MibBuilder().loadModules()
      self._view = view.MibViewController(self._builder)
      self._dispatch = AsynsockDispatcher()
      self._started = time()
      self._retval = {}

   _getid = property(fget = lambda self: self._snmp.apiPDU.getRequestID)
   value = property(fget = lambda self: self._retval)
   close = property(fget = lambda self: self._dispatch.closeDispatcher())
   integer = property(fget = lambda self: self._snmp.Integer)
   integer32 = property(fget = lambda self: self._snmp.Integer32)
   ipaddress = property(fget = lambda self: self._snmp.IpAddress)
   null = property(fget = lambda self: self._snmp.Null)
   string = property(fget = lambda self: self._snmp.OctetString)
   timer = property(fget = lambda self: self._snmp.TimeTicks)
   unsigned = property(fget = lambda self: self._snmp.Unsigned32)
   opaque = property(fget = lambda self: self._snmp.Opaque)

   def _error(self, Message):
      if Message:
         raise SNMPError("%s: %s" %(str(self), str(Message)))

   def __str__(self):
      return "%s: %d" %(self._host, self._port)

   def __repr__(self):
      return repr(str(self))

   def __unicode__(self):
      return unicode(str(self))

   def _setpdu(self, snmptype):
      if snmptype == GET:
         self._reqpdu = self._snmp.GetRequestPDU()
      elif snmptype == SET:
         self._reqpdu = self._snmp.SetRequestPDU()
      elif snmptype == WALK:
         self._reqpdu = self._snmp.GetNextRequestPDU()

      self._snmp.apiPDU.setDefaults(self._reqpdu)
      self._type = snmptype

   def _timer(self, now, started = time()):
      if (now - started) > 3:
         self._error('request timed out!')

   def _transport(self):
      if not self._conn:
         self._reqmesg = self._snmp.Message()
         self._snmp.apiMessage.setDefaults(self._reqmesg)
         self._snmp.apiMessage.setCommunity(self._reqmesg, self._comm)
         self._conn = True

      self._snmp.apiMessage.setPDU(self._reqmesg, self._reqpdu)
      self._dispatch.registerTransport(udp.domainName,
            udp.UdpSocketTransport().openClientMode())
      self._dispatch.registerRecvCbFun(self._recvCallBack)
      self._dispatch.sendMessage(encoder.encode(self._reqmesg),
            udp.domainName, (self._host, self._port))
      self._dispatch.jobStarted(True)
      self._dispatch.runDispatcher()
      self._dispatch.closeDispatcher()

   def _recvCallBack(self, dispatcher, domain, address, message, reqpdu = None):
      if not reqpdu:
         reqpdu = self._reqpdu

      while message:
         response, message = decoder.decode(message,
               asn1Spec = self._snmp.Message())
         respdu = self._snmp.apiMessage.getPDU(response)
         if self._getid(reqpdu) == self._getid(respdu):
            status = self._snmp.apiPDU.getErrorStatus(respdu)
            if (self._type == GET) or (self._type == SET):
               if status:
                  self._error(status.prettyPrint())

               for oid, val in self._snmp.apiPDU.getVarBinds(respdu):
                  self._retval[oid.prettyPrint()] = val.prettyPrint()
            
               self._dispatch.jobFinished(True)
            elif self._type == WALK:
               if status and (status != 2):
                  self._error(status.prettyPrint())

               varbind = self._snmp.apiPDU.getVarBindTable(reqpdu, respdu)
               for row in varbind:
                  for name, val in row:
                     if val is None:
                        continue

                     self._retval[name.prettyPrint()] = val.prettyPrint()

               for oid, val in varbind[-1]:
                  if val is not None:
                     break

               else:
                  self._dispatch.jobFinished(True)

               self._snmp.apiPDU.setVarBinds(
                     reqpdu, map(lambda (ix, jy), val = self._snmp.Null(): (ix,
                        val), varbind[-1]))

               self._snmp.apiPDU.setRequestID(reqpdu,
                     self._snmp.getNextRequestID())
               self._dispatch.sendMessage(
                     encoder.encode(self._reqmesg), domain, address)

               if (time() - self._started) > 3:
                  self._error('request timed out')

               self._started = time()

      return message

   def Get(self, Oid):
      self._setpdu(GET)
      self._snmp.apiPDU.setVarBinds(
            self._reqpdu, ((Oid, self._snmp.Null()),))
      self._transport()
      return self._retval.values()

   def Set(self, Oid, Value):
      self._setpdu(SET)
      self._snmp.apiPDU.setVarBinds(self._reqpdu, ((Oid, Value), ))
      self._transport()
      return self._retval.values()

   def Walk(self, Oid):
      self._setpdu(WALK)
      headVars = [ self._snmp.ObjectIdentifier((Oid)) ]
      self._snmp.apiPDU.setVarBinds(self._reqpdu, map(
         lambda ix, snmp = self._snmp: (ix, self._snmp.Null()), headVars))
      self._transport()
      return self._retval


snmpv3auth = stdlib.EnumType('SHA', 'MD5')
snmpv3enc = stdlib.EnumType('AES', 'AES-192', 'AES-256', 'TDES', 'DES')

class snmpv3(object):
   def __init__(self, Host, Port = 161, **Keywords):
      self._user = None
      self._authpass = None
      self._encpass = None
      self._auth = None
      self._enc = None
      self._host = None
      self._port = 0
      self._cbctx = {}
      self._retval = {}
      self._snmptype = None

      self._engine = engine.SnmpEngine()
      self._param = stdlib.GetRandomString(6)
      self._target = stdlib.GetRandomString(6)

      if Host:
         self._host = Host
      else:
         raise SNMPError("No hostname defined!")

      if Port:
         try:
            self._port = int(Port)
         except ValueError as err:
            raise SNMPError(err.message)
      else:
         self._port = 161

      if "user" in Keywords:
         self._user = Keywords["user"].strip()
      else:
         self._error("No username defined for snmp session!")

      if "authpass" in Keywords:
         self._authpass = Keywords["authpass"].strip()

      if "privpass" in Keywords:
         self._privpass = Keywords["privpass"].strip()

      if "auth" in Keywords:
         self._auth = Keywords["auth"]

      if "enc" in Keywords:
         self._enc = Keywords["enc"]

      self._aflag = True
      self._pflag = True

      self._authproto = None
      if snmpv3auth[self._auth] == "SHA":
         self._authproto = config.usmHMACSHAAuthProtocol
      elif snmpv3auth[self._auth] == "MD5":
         self._authproto = config.usmHMACMD5AuthProtocol
      else:
         self._aflag = False
         self._authproto = config.usmNoAuthProtocol

      self._encproto = None
      if snmpv3enc[self._enc] == "AES":
         self._encproto = config.usmAesCfb128Protocol
      elif snmpv3enc[self._enc] == "AES-192":
         self._encproto = config.usmAesCfb192Protocol
      elif snmpv3enc[self._enc] == "AES-256":
         self._encproto = config.usmAesCfb256Protocol
      elif snmpv3enc[self._enc] == "TDES":
         self._encproto = config.usm3DESEDEPrivProtocol
      elif snmpv3enc[self._enc] == "DES":
         self._encproto = config.usmDESPrivProtocol
      else:
         self._pflag = False
         self._encproto = config.usmNoPrivProtocol

      self._type = None
      if self._aflag and self._pflag:
         self._type = 'authPriv'
      elif self._aflag and not self._pflag:
         self._type = 'authNoPriv'
      elif not self._aflag and not self._pflag:
         self._type = 'noAuthNoPriv'

      config.addV3User(self._engine, self._user,
            self._authproto, self._authpass, self._encproto, self._privpass)
      config.addTargetParams(self._engine, self._param, self._user,
            self._type)
      config.addTargetAddr(self._engine, self._target,
            config.snmpUDPDomain, (self._host, self._port), self._param)
      config.addSocketTransport(self._engine, udp.domainName,
            udp.UdpSocketTransport().openClientMode())

   retval = property(fget = lambda self: self._retval)
   """
   (1,3,6,1,2,1,1,1,0), rfc1902.OctetString('Grinch')
   """

   @property
   def OctetString(self):
      return rfc1902.OctetString

   @property
   def Counter32(self):
      return rfc1902.Counter32

   @property
   def Counter64(self):
      return rfc1902.Counter64

   @property
   def Gauge32(self):
      return rfc1902.Gauge32

   @property
   def Integer(self):
      return rfc1902.Integer

   @property
   def Integer32(self):
      return rfc1902.Integer32

   @property
   def IPAddress(self):
      return rfc1902.IpAddress

   IpAddress = property(fget = lambda self: self.IPAddress)

   @property
   def OctetString(self):
      return rfc1902.OctetString

   @property
   def Unsigned32(self):
      return rfc1902.Unsigned32

   @property
   def Opaque(self):
      return rfc1902.Opaque

   @property
   def TimeTicks(self):
      return rfc1902.TimeTicks

   def _callback(self, sendRequestHandle, errorIndication, errorStatus,
         errorIndex, varBinds, cbctx):
      retval = None
      if self._snmptype in (snmptype.GET, snmptype.SET):
         cbctx['errorIndication'] = errorIndication
         cbctx['errorStatus'] = errorStatus
         cbctx['errorIndex'] = errorIndex
         cbctx['varBinds'] = varBinds
         cbctx['handle'] = sendRequestHandle
         self._cbctx = cbctx

         if errorIndication:
            self._error("snmp error indication for %s: %s", self._host,
                  errorIndication)
         elif errorStatus:
            if errorIndex:
               self._error("%s: %s at %s", self._host,
                     errorStatus.prettyPrint(), varBinds[int(errorIndex) - 1])
            else:
               self._error("snmp error status for %s: %s", self._host,
                     errorStatus.prettyPrint())
         else:
            for (oid, val) in varBinds:
               self._retval[oid.prettyPrint()] = val.prettyPrint()
      elif self._snmptype == snmptype.WALK:
         if errorIndication:
            self._error("%s: %s", self._host, errorIndication)
         
         if errorStatus and errorStatus != 2:
            self._error("%s: %s at %s", self._host,
                  errorStatus.prettyPrint(),
                  errorIndex and varBinds[-1][int(errorIndex)-1] or '?')

         for varBindRow in varBinds:
            for oid, val in varBindRow:
               self._retval[oid.prettyPrint()] = val.prettyPrint()

         retval = 1

      return retval

   def _error(self, fmt, *args):
      mesg = stdlib.sprintf(fmt, *args)
      raise SNMPError(mesg)

   def Get(self, Oid):
      self._snmptype = snmptype.GET
      cmdgen.GetCommandGenerator().sendReq(self._engine, self._target,
            ((Oid, None),), self._callback, self._cbctx)
      self._engine.transportDispatcher.runDispatcher()
      return self._retval

   def Set(self, Oid, Value):
      self._snmptype = snmptype.SET
      cmdgen.SetCommandGenerator().sendReq(self._engine, self._target,
            ((Oid, Value),), self._callback, self._cbctx)
      self._engine.transportDispatcher.runDispatcher()
      return self._retval

   def Walk(self, Oid):
      self._snmptype = snmptype.WALK
      cmdgen.NextCommandGenerator().sendReq(
            self._engine, self._target, ((Oid, None),), self._callback)
      self._engine.transportDispatcher.runDispatcher()
      return self._retval
